from configure import *
import tlsh
import os
from tqdm import tqdm
import sys
sys.path.append("/home/cyw/projects/function_sim_project/basic_script")


class mytlsh():

    def modify_res(self, res):
        """
            将tlsh的score 转换成0到1之间
            根据推荐的指标100为临界点--》0.5
            100的正好为0.5 ，最大值=1，最小值等于0（样本量较少的话，可能会有bug）
        """
        # a=max(res)
        # b=min(res)
        # temp=[0]*len(res)
        # for i in range(len(res)):
        #     b_min=max(b,100)
        #     a_max=min(a,100)
        #     if res[i] > 100:
        #         tar=(1-(res[i]-b_min)/(a-b_min))/2
        #     else:
        #         tar=(1-(res[i]-b)/(a_max-b))/2+0.5
        #     temp[i]=tar
        a = max(res)
        temp = [0]*len(res)
        for i in range(len(res)):
            if res[i] > 100:
                tar = (1 - res[i]/a)/2
            else:
                tar = (1 - res[i]/(100))/2 + 0.5
            temp[i] = tar
        return temp

    def score(self, a, b):
        """
            输入:对比的两个样本的名字
            输出:这两个样本的tlsh得分
            这里设置的推荐阈值是100
            score越低--->相似度越高
        """
        res = []
        for i in range(len(a)):
            # dataPath = r"/mnt/mydisk1/chenyongwei/data/functionsim/allorigindata/"
            dataPath = r"/mnt/mydisk1/chenyongwei/malware/BODMAS_dataset/dataTrans/"
            h1 = tlsh.hash(open(dataPath+a[i], 'rb').read())
            h2 = tlsh.hash(open(dataPath+b[i], 'rb').read())
            score = tlsh.diff(h1, h2)
            res.append(score)
        return self.modify_res(res)


if __name__ == "__main__":
    a = mytlsh()
    # score=a.score(["cf0d39927cc9c7345f5c5f97dc78ef7c","cf0d39927cc9c7345f5c5f97dc78ef7c","0b9a015228b264f23d77fee0c098a9e1"],["cf0d39927cc9c7345f5c5f97dc78ef7c","ff85314524c696fbee2d6288d594ac12","cf0d39927cc9c7345f5c5f97dc78ef7c"])
    score = a.score(["cf0d39927cc9c7345f5c5f97dc78ef7c", "0b9a015228b264f23d77fee0c098a9e1",], [
                    "cf0d39927cc9c7345f5c5f97dc78ef7c", "ff85314524c696fbee2d6288d594ac12"])
    print(score)
